import os
import pandas as pd


def distinct(list_of_hypotheses, n):
    count = 0
    total_score = 0.0

    for multi_hypotheses in list_of_hypotheses:
        # drop the index number
        # for idx, hyp in enumerate(multi_hypotheses):
        #     if hyp == "":
        #         continue
        #     try:
        #         multi_hypotheses[idx] = hyp[hyp.index(']') + 1:]
        #     except:
        #         # print(hyp)
        #         pass
        multi_hypotheses = [hyp.split() for hyp in multi_hypotheses]
        ngram_set = set()
        ngram_sum = 0
        for hyp in multi_hypotheses:
            length = len(hyp) - n + 1
            for i in range(length):
                ngram_sum += 1
                ngram_set.add(' '.join(hyp[i:i + n]))
        if ngram_sum > 0:
            count += 1
            total_score += (len(ngram_set) / ngram_sum)

    return total_score / count if count > 0 else 0


def cal_distinct_n(list_of_hypotheses):
    def r(x):
        return round(x, 4)

    dist_1 = distinct(list_of_hypotheses, 1)
    dist_2 = distinct(list_of_hypotheses, 2)
    dist_3 = distinct(list_of_hypotheses, 3)
    dist_4 = distinct(list_of_hypotheses, 4)
    print('dist-1: {:.4f}\tdist-2: {:.4f}\tdist-3: {:.4f}\tdist-4: {:.4f}'.format(r(dist_1), r(dist_2), r(dist_3),
                                                                                  r(dist_4)))


# pass the filename and then return the distinct metric
def get_distinct_by_file(predition_file):
    evaluation_predictions = pd.read_csv(predition_file)
    # print(evaluation_predictions[evaluation_predictions.isnull().values == True])
    evaluation_predictions.fillna("", inplace=True)
    evaluation_predictions3 = evaluation_predictions.iloc[:, 1:4].values.tolist()
    evaluation_predictions5 = evaluation_predictions.iloc[:, 1:6].values.tolist()
    evaluation_predictions10 = evaluation_predictions.iloc[:, 1:11].values.tolist()

    cal_distinct_n(evaluation_predictions3)
    cal_distinct_n(evaluation_predictions5)
    cal_distinct_n(evaluation_predictions10)


# calculate single file
prediction_file = '../output/predictions/t5-small-title-generation-lowercase-alldata-bs8-402-23/predictions-bs.csv'
get_distinct_by_file(prediction_file)

# calculate multi file in one directory
# prediction_dir = '../output/predictions/t5-small-title-generation-402-23/'
# files = os.listdir(prediction_dir)
# files.sort()
# for file in files:
#     print("\n" + file)
#     get_distinct_by_file(prediction_dir + file)


